#include "websocketserver.h"
#include <QDebug>
#include <iostream>

// WebSocket会话构造函数：接收已连接的socket
WebSocketSession::WebSocketSession(tcp::socket socket)
    : ws_(std::move(socket)), closed_(false) { // 初始化WebSocket流和关闭状态
}

// 启动会话：开始接受客户端连接
void WebSocketSession::start() {
    // 设置超时和响应头
    ws_.set_option(websocket::stream_base::timeout::suggested(beast::role_type::server));
    ws_.set_option(websocket::stream_base::decorator([](websocket::response_type& res) {
        res.set(http::field::server, "VideoChat WebSocket Server");
        res.set(http::field::access_control_allow_origin, "*");  // 允许跨域
    }));

    // 异步接受WebSocket握手
    ws_.async_accept([self = shared_from_this()](beast::error_code ec) {
        if (!ec) {
            self->doRead();  // 接受成功后开始读取数据
        } else {
            qDebug() << "WebSocket accept error:" << ec.message().c_str();
            self->close();  // 出错则关闭
        }
    });
}

// 发送消息：支持文本和二进制类型
void WebSocketSession::send(const std::string& message, bool is_binary) {
    if (closed_.load()) return;  // 如果已关闭则不发送

    // 设置帧类型（文本/二进制）
    ws_.binary(is_binary);

    // 异步发送消息
    ws_.async_write(asio::buffer(message),
                    [self = shared_from_this()](beast::error_code ec, std::size_t) {
                        if (ec) {  // 发送失败处理
                            qDebug() << "WebSocket send error:" << ec.message().c_str();
                            self->close();
                        }
                    });
}

// 关闭会话
void WebSocketSession::close() {
    if (closed_.exchange(true)) return;  // 原子操作检查并设置关闭状态

    if (ws_.is_open()) {  // 如果连接仍打开
        boost::system::error_code ec;
        // 异步关闭WebSocket连接
        ws_.async_close(websocket::close_code::normal,
                        [self = shared_from_this()](beast::error_code ec) {
                            boost::system::error_code ignored_ec;
                            // 关闭底层socket
                            self->ws_.next_layer().shutdown(tcp::socket::shutdown_both, ignored_ec);
                            self->ws_.next_layer().close(ignored_ec);

                            // 调用关闭回调
                            if (self->closeHandler_) {
                                self->closeHandler_(self);
                            }
                        });
    } else {  // 如果已关闭，直接调用回调
        if (closeHandler_) {
            closeHandler_(shared_from_this());
        }
    }
}

// 设置消息处理回调
void WebSocketSession::setMessageHandler(MessageHandler handler) {
    messageHandler_ = handler;
}

// 设置关闭处理回调
void WebSocketSession::setCloseHandler(CloseHandler handler) {
    closeHandler_ = handler;
}

// 开始异步读取数据
void WebSocketSession::doRead() {
    if (closed_.load()) return;  // 已关闭则不读取

    // 异步读取消息
    ws_.async_read(buffer_,
                   [self = shared_from_this()](beast::error_code ec, std::size_t bytes_transferred) {
                       self->onRead(ec, bytes_transferred);  // 读取完成回调
                   });
}

// 读取完成处理
void WebSocketSession::onRead(beast::error_code ec, std::size_t bytes_transferred) {
    if (ec) {  // 读取错误处理
        if (ec != websocket::error::closed) {  // 非正常关闭则输出错误
            qDebug() << "WebSocket read error:" << ec.message().c_str();
        }
        close();  // 关闭会话
        return;
    }

    if (messageHandler_) {
        // 如果设置了消息处理器
        bool is_binary = ws_.got_binary();  // 判断是否为二进制消息
        std::string message = beast::buffers_to_string(buffer_.data());  // 转换为字符串
        messageHandler_(shared_from_this(), message, is_binary);  // 调用回调
    }

    buffer_.consume(bytes_transferred);  // 清理缓冲区

    if (!closed_.load()) {  // 继续读取下一条消息
        doRead();
    }
}

// WebSocket服务器构造函数
WebSocketServer::WebSocketServer(QObject *parent)
    : QObject(parent), ioc_(), acceptor_(ioc_) {
}

// 服务器析构函数：停止服务器
WebSocketServer::~WebSocketServer() {
    stop();
}

// 启动服务器：指定端口
bool WebSocketServer::start(quint16 port) {
    if (running_.load()) return true;  // 已运行则直接返回

    try {
        tcp::endpoint endpoint(tcp::v4(), port);  // 创建端点（IPv4 + 端口）
        acceptor_.open(endpoint.protocol());      // 打开 acceptor
        acceptor_.set_option(asio::socket_base::reuse_address(true));  // 允许地址重用
        acceptor_.bind(endpoint);                 // 绑定端点
        acceptor_.listen(asio::socket_base::max_listen_connections);   // 开始监听

        running_.store(true);  // 设置运行状态
        // 启动IO线程
        ioThread_ = std::thread([this]() {
            doAccept();  // 开始接受连接
            ioc_.run();  // 运行IO服务
        });

        qDebug() << "WebSocket server started on port" << port;
        return true;
    } catch (const std::exception& e) {  // 启动失败处理
        qDebug() << "Failed to start WebSocket server:" << e.what();
        return false;
    }
}

// 停止服务器
void WebSocketServer::stop() {
    if (!running_.load()) return;  // 未运行则直接返回

    running_.store(false);  // 设置停止状态

    // 先关闭acceptor
    boost::system::error_code ec;
    acceptor_.close(ec);

    // 然后关闭所有会话
    {
        std::lock_guard<std::mutex> lock(sessionsMutex_);  // 加锁保护会话集合
        for (auto& session : sessions_) {
            session->close();
        }
        sessions_.clear();
    }

    // 停止IO上下文
    ioc_.stop();

    // 等待IO线程结束
    if (ioThread_.joinable()) {
        ioThread_.join();
    }

    ioc_.restart();  // 重置IO上下文
    qDebug() << "WebSocket server stopped";
}

// 广播消息给所有客户端
void WebSocketServer::broadcast(const std::string& message, bool is_binary) {
    std::lock_guard<std::mutex> lock(sessionsMutex_);  // 加锁保护会话集合

    // 清理已关闭的会话并广播消息
    for (auto it = sessions_.begin(); it != sessions_.end(); ) {
        if ((*it)->isClosed()) {  // 移除已关闭的会话
            it = sessions_.erase(it);
        } else {
            (*it)->send(message, is_binary);  // 发送消息
            ++it;
        }
    }
}

// 广播消息给所有客户端（排除指定会话）
void WebSocketServer::broadcast(const std::string& message, bool is_binary, std::shared_ptr<WebSocketSession> excludeSession) {
    std::lock_guard<std::mutex> lock(sessionsMutex_);  // 加锁保护会话集合

    // 清理已关闭的会话并广播消息
    for (auto it = sessions_.begin(); it != sessions_.end(); ) {
        if ((*it)->isClosed()) {  // 移除已关闭的会话
            it = sessions_.erase(it);
        } else {
            // 如果不是排除的会话，则发送消息
            if (*it != excludeSession) {
                (*it)->send(message, is_binary);
            }
            ++it;
        }
    }
}

// 移除会话
void WebSocketServer::removeSession(std::shared_ptr<WebSocketSession> session) {
    std::lock_guard<std::mutex> lock(sessionsMutex_);  // 加锁保护会话集合
    sessions_.erase(session);
    emit clientDisconnected();  // 发射客户端断开信号
}

// 开始接受新连接
void WebSocketServer::doAccept() {
    if (!running_.load()) return;  // 未运行则不接受

    // 异步接受连接
    acceptor_.async_accept([this](beast::error_code ec, tcp::socket socket) {
        if (!ec) {  // 接受成功
            // 创建新会话
            auto session = std::make_shared<WebSocketSession>(std::move(socket));

            // 设置消息处理器
            session->setMessageHandler([this](std::shared_ptr<WebSocketSession> session,
                                              const std::string& message,
                                              bool is_binary) {
                if (is_binary) {  // 二进制消息（视频帧）
                    broadcast(message, true, session);  // 广播给其他客户端
                    emit binaryMessageReceived(QByteArray::fromStdString(message));
                } else {  // 文本消息
                    emit messageReceived(QString::fromStdString(message));  // 发射文本消息信号
                    broadcast(message, false, session);  // 广播给其他客户端
                }
            });

            // 设置关闭处理器
            session->setCloseHandler([this](auto session) {
                removeSession(session);  // 移除会话
            });

            // 将新会话添加到集合
            {
                std::lock_guard<std::mutex> lock(sessionsMutex_);
                sessions_.insert(session);
            }

            session->start();  // 启动会话
            emit clientConnected();  // 发射客户端连接信号
        } else {  // 接受失败
            qDebug() << "Accept error:" << ec.message().c_str();
        }

        if (running_.load()) {  // 继续接受下一个连接
            doAccept();
        }
    });
}
